import numpy as np, scipy.special as sp, scipy.integrate as ig, copy, re, math
import sys
from Model import Model,Actor,Evidence,TrackSegment,Node,TimeNode,Time,Distributions

###
# Use null models to calculate ratio of conditional density to prior density for partially explained models, then use this ratio to adjust some basic prior for unobserved nodes and multiply by a global prior for the model itself to hypothesize a likelihood given incomplete evidence


# Using conditional densities for the node's parameters, we want the expected score value for each one.  We can get this by computing, for a parameter B for which we have a distribution conditioned on node l f(b) = p(B=b|l)  E(f(B)) = sum_p{f(b)*p(l)*p(b|l) + f(b)*p(~l)*p(b|~l)} = E(f(B)|l)+E(f(B)|~l) = p(l)*sum_b{f(b)*p(b|l)} + p(~l)*sum_p{f(b)*p(b|~l)}

# ***Don't need p(~l) term; we want the conditional expectation of the score since the model returns the conditional score to the blackboard
# We need to calculate E(f(B)|l) = p(l)*sum_b{f(b)*p(b|l)}***

# ***Also, we don't want E(f(B)|l) for each piece of evidence individually; we want the expected joint density across all nodes.
# The nodes are still independent conditioned on l, so we just calculate their conditional expectations individually (E(f(B_i)|l) = sum_b{f(b)*p(b|l)}), then take that product multiplied by l's prior and by the product of the densities for any observed evidence (if some of our density functions can be calculated without all the evidence being observed)***

# However, while we have some initial value for p(l), we can weight this by our belief in other nodes which we have observed.  Suppose node n has a conditional density value of a1(n), while its prior model has a density of a2(n).  We can use p'(l) = a1(n)/a2(n)*p(l), and p*(l) = p'(l)/(p'(l)+p(~l)) 
###

###
# Prior on lag time score: E(f(L)|R2) = sum_l{f(l)*P(l|R2)} (expectation on lag time score conditioned on run2, given by summing over possible lag time values multiplied by their conditional probability
###


class Chase:

    def __init__(self):
        """Chase class has an internal model which contains all of our chase-specific domain knowledge"""
        self.model = Model()
        self.model.name = "Chase"
        self.model.subscriptions = ["MoveFast"]
        """subscriptions currently don't do anything, but this is how the model will know which types of evidence to subscribe to on the blackboard"""
        self.model.topNode = "chasetop"
        self.model.independent = True
        # self.model.subjects = []
        # self.model.objects = []
        # self.model.ind_objects = []

        ###
        # create placeholder actors in order to build the chase template (these are what are used to match newly received evidence's actors across the entire model instead of just at the node corresponding to the specific piece of evidence)
        ###
        act1 = Actor()
        act1.name = "dummy1"
        act1.atype = Actor.NULL
        act2 = Actor()
        act2.name = "dummy2"
        act2.atype = Actor.NULL

        # self.model.subjects.append(act1)
        # self.model.objects.append(act2)


        """chase model has 3 nodes: the top level activity node, a node which represents whether the object (chasee) is running, and a node which represents whether the subject (chaser) is both running and chasing the object"""

        """top level node template"""
        tnode = Node()
        tnode.getID()
        tnode.name = "chasetop"
        tnode.subjects = [act1]
        tnode.objects = [act2]
        tnode.evvecs = []
        tnode.bound = False
        pnode = copy.deepcopy(tnode)
        self.model.topNode = pnode.name
        self.model.nodes.append(pnode)

        """chasee's run template"""
        tnode.name = "run1"
        tnode.getID()
        tnode.subjects = [act2]
        tnode.objects = []
        tnode.parents = [pnode]
        tnode.hasPrior = True
        tnode.evCard = 1.
        tnode.reqActs = dict()
        tnode.reqActs["subject"] = 1
        tnode.atEnd = False
        tnode.bound = False
        """just explains act2's running speed"""
        #tnode.prior = np.log(.05)
        """latent nodes in the chase model have the top-level node as a parent"""
        
        """template evidence object used to find valid bindings for evidence we receive"""
        tev = TrackSegment()
        tev.actor = act2
        tev.segtype = Evidence.MOVE_FAST
        tev.willBind = True
        tev.endsNode = True
        tev.bound = False
        """evidence labels are used to match evidence we receive from the blackboard to nodes which make predictions about that type of evidence (the run node only predicts parameters of tracks, and only cares about segments of track when the actor seems to be moving at a faster pace than walking)"""
        tnode.evvecs = [copy.deepcopy(tev)]

        """This function corresponds to the distribution of a runner's speed conditioned on absence of a chase (i.e. default running speed)"""
        def moveNull1(evvecs):
            """runner's location vector"""
            t = evvecs[0]
            locs = t.locations
            dt = float(len(locs))
            speed = t.calculate_speed()
            weight = Chase.speedDensityNull(speed)
            return weight

        """each node has a likelihood function which calculates the joint conditional likelihood for all of the evidence it's predicting (this likelihood unique to a node's subject, object, indirect object, and time bindings)"""
        def moveJoint1(evvecs):
            t = evvecs[0]
            """location vector for runner is only item in evvec"""
            locs = t.locations
            dt = float(len(locs))
            speed = t.calculate_speed()
            tweight = Chase.speedDensityChase(speed)
            """a single person's run likelihood is only based on his movement speed"""
            fweight = moveNull1(evvecs)
            return tweight,fweight
        tnode.lfunction = moveJoint1
        self.model.nodes.append(copy.deepcopy(tnode))


        """template node that predicts parameters of the chaser's run"""
        tnode.name = "run2"
        tnode.getID()
        tnode.subjects = [act1]
        tnode.objects = [act2]
        tnode.parents = [pnode]
        tnode.evCard = 3.
        tnode.bound = False
        tnode.atEnd = True
        """for both actors, explains their run speed, then explains their relative angle"""

        """chaser's run node has two evidence segments, since it uses the chasee's run direction and start time to calculate the likelihood of the chaser's run direction and start time"""
        tev.actor = act1
        tev.segtype = Evidence.MOVE_FAST
        tev.timeInd = 1
        tev.willBind = True
        tev.endsNode = True
        tev.bound = False
        tnode.evvecs = [copy.deepcopy(tev)]

        tev.actor = act2
        tev.segtype = Evidence.MOVE_FAST
        tev.timeInd = 0
        tev.willBind = True
        tev.endsNode = False
        tnode.hasPrior = True
        tev.bound = False
        #tnode.reqActs = [act2] #don't have prior on first person given second
        #tnode.prior = np.log(.02)
        tnode.reqActs = dict()
        tnode.reqActs["object"] = 1
        tnode.evvecs.append(copy.deepcopy(tev))        
        
        """this run likelihood function calculates the likelihood that the chaser is running, that he is moving in the same direction as the chasee, and that the lag time between them indicates a chase (these values are conditionally independent given that the chaser is chasing the chasee, so their joint likelihood is just their product)"""
        def moveJoint2(evvecs):
            t1 = evvecs[0]
            t2 = evvecs[1]
            locs = t1.locations
            olocs = t2.locations
            """locs: subject's location vector, olocs: object's location vector"""
            speed = t1.calculate_speed()
            sweight = Chase.speedDensityChase(speed)
            """sweight is a conditional density based on subject's speed"""
            angle = t1.calculate_direction()
            oang = t2.calculate_direction()
            aweight = Chase.angDensityChase(angle,oang)
            """aweight is a conditional density based on the relative angle between the subject's and object's tracks"""
            #tweight = Chase.timeDensity(tbet,endtime-ost)
            """tweight is a conditional density based on the time between the chasee's appearance and the chaser's appearance"""
            tjoint = sweight + aweight
            fjoint = moveNull2(evvecs)
            #joint = sweight + aweight + tweight
            #print "sweight: "+str(np.exp(sweight))+", aweight: "+str(np.exp(aweight))+", tweight: "+str(np.exp(tweight))
            """evidence likelihoods are (locally) independent conditioned on latent node,
            so their joint is just their sum (log values)"""
            return tjoint,fjoint

        def moveNull2(evvecs):
            t1 = evvecs[0]
            t2 = evvecs[1]
            locs1 = t1.locations
            locs2 = t2.locations
            speed = t1.calculate_speed()
            sweight = Chase.speedDensityNull(speed)
            ang = t1.calculate_direction()
            oang = t2.calculate_direction()
            aweight = Chase.angDensityNull(ang,oang)
            joint = sweight + aweight
            return joint

        def eR2Score(evvecs):
            t = evvecs[0]
            """We're in the expected score function, so we know that the chasee's evidence has been observed and the chaser's has not"""
            eSScore,ignore = ig.quad(Chase.sqSpeedDensity,0,float("inf"))
            """speed, so defined on the interval [0,+inf)"""
            # print "calculated expected score for speed with subject "+evvecs[0].actor.name+" and object "+evvecs[1].actor.name+"; score is "+str(np.log(eSScore))
            # print "extra information returned:"
            # print ignore
            eAScore,ignore = ig.quad(Chase.sqAngleDensity,0,1)
            """relative angle is defined on the interval [-pi,pi], but beta distributions are defined on [0,1], so when calculating the angle's density we scale it into that range.  Not necessary here since the integral should be the same."""
            # print "calculated expected score for relative angle with subject "+evvecs[0].actor.name+" and object "+evvecs[1].actor.name+"; score is "+str(np.log(eAScore))
            # print "extra information returned:"
            # print ignore
            if (eSScore <= 0) or (eAScore <= 0):
                return float("-inf")
            else:
                weight = np.log(eSScore) + np.log(eAScore)
            return weight

        tnode.lfunction = moveJoint2
        tnode.eDen = eR2Score
        self.model.nodes.append(copy.deepcopy(tnode))
        pnode = tnode

        tnode = TimeNode()
        tnode.name = "lagtime"
        tnode.getID()
        tnode.subjects = [act2]
        tnode.objects = [act1]
        tnode.reqActs = dict()
        tnode.reqActs["subject"] = 1
        tnode.lagfunction = Chase.compoundTimeDensity
        tnode.parents = [pnode]
        tnode.evCard = 1.
        tnode.atEnd = True
        tnode.bound = False
        """only counting as 1 piece of evidence, since we're not actually predicting either person's arrival time individually"""

        tev.actor = act2
        tev.segtype = Evidence.MOVE_FAST
        tev.timeInd = 0
        tev.willBind = True
        tev.atEnd = False
        tev.bound = False
        tnode.evvecs = [copy.deepcopy(tev)]

        tev.actor = act1
        tev.segtype = Evidence.MOVE_FAST
        tev.timeInd = 0
        tev.willBind = True
        tev.endsNode = True
        tev.bound = False
        tnode.evvecs.append(copy.deepcopy(tev))

        tev = Time()
        tev.segtype = Evidence.WAKE
        tev.subjects = []
        tev.objects = []
        tev.timeInd = 0
        tev.willBind = False
        tnode.evvecs.append(copy.deepcopy(tev))
        tnode.hasPrior = True
        #tnode.reqActs = [act2] #don't have prior on lag time given only the second person
        #tnode.prior = np.log(.05)

        def tDensity(evvecs,ctime = None):
            if not evvecs[0].actor.bound:
                """can't say anything meaningful about lag time when first person is unobserved; shouldn't ever get called in this state anyway"""
                return float("nan"),float("nan")
            if evvecs[1].actor.bound:
                #print "have both actors?"
                """prefer arrival time prediction to time waited"""
                tdif = float(evvecs[1].starttime-evvecs[0].endtime)/Model.FR
                #print str(evvecs[0].starttime)+" "+str(evvecs[1].starttime)
                #print "timedensity "+str(Chase.timeDensity(tdif))+"\n"
                return Chase.timeDensityChase(tdif),.5
            if not ctime is None:
                tdif = float(ctime-evvecs[0].endtime)/Model.FR
                #print "waitdensity ctime "+str(Chase.waitDensity(tdif))+"\n"
                tval = Chase.waitDensity(tdif)
                return tval,1-tval,.5
            if evvecs[2].endtime != -1:
                tdif = float(evvecs[2].endtime-evvecs[0].endtime)/Model.FR
                #print "waitdensity "+str(Chase.waitDensity(tdif))+"\n"
                return Chase.waitDensity(tdif),.5
            """can't say anything meaningful about lag time without knowing either arrival time or current time"""
            #print "didn't find end time????\n"
            return float("nan"),float("nan")

        tnode.lfunction = tDensity

        def eTScore(evvecs,cTime = None):
            tDif = 0
            #print "in eTScore, with evvecs:"
            #print "\n".join(["name: "+x.actor.name+", starttime: "+str(x.starttime)+", endtime: "+str(x.endtime)+", bound: "+str(x.bound) for x in evvecs])
            if not cTime is None:
                tDif = float(cTime - evvecs[0].endtime)/Model.FR
            elif evvecs[2].endtime != -1:
                tDif = float(evvecs[2].endtime - evvecs[0].endtime)/Model.FR
            """Same as in eR2Score, we know evvecs[0] is present (and evvecs[1] is not) since we're in the expectation function"""
            eLScore,ignore = ig.quad(Chase.sqTimeDensity,tDif,float("inf"))
            # print "calculated expected score for lag time with subject "+evvecs[0].actor.name+" and object "+evvecs[1].actor.name+"; score is "+str(np.log(eLScore))
            # print "extra information returned:"
            # print ignore
            """Lag time is defined on the interval [0,+inf), but we know that we can't have a lag time less than the time that has already elapsed"""
            if eLScore <= 0:
                return float("-inf")
            return np.log(eLScore)

        tnode.eDen = eTScore
        timenode = copy.deepcopy(tnode)
        self.model.nodes.append(timenode)
        self.model.timeNodes[timenode.name] = timenode

    """these functions just delegate calls to the chase class's internal model
    this should probably be refactored at some point"""
    def genXML(self):
        return self.model.generate_xml()

    def printXML(self,outf = None,mfiles = False):
        return self.model.print_xml(outf,mfiles)

    # def addEv(self,ev):
    #     #print "adding evidence"
    #     self.model.addEv(ev)

    # def evalEv(self):
    #     self.model.evalEv()

    """the methods below call the PDF functions defined in models.py but use parameters which are specific to the chase model"""

    @staticmethod
    def compoundTimeDensity(timediff,arrived = False):
        if arrived:
            return timeDensity(timediff)
        return waitDensity(timediff)

    @staticmethod
    def waitDensity(timediff):
        if timediff < 0:
            return float("-inf")
        tweight = Distributions.Surv.gamma(timediff,1.5,1./2.)
        return tweight
    
    @staticmethod
    def timeDensityChase(timediff):
        if timediff < 0:
            """gamma dist isn't defined for negative values, and in this chase model the chaser appears after the chasee anyway"""
            return float("-inf")
        weight = Distributions.PDFs.gamma(timediff,3.75,1.5)
        """estimate time between 1st and 2nd person drawn from gamma with alpha = 3.75, beta = 1.5"""
        return weight

    @staticmethod
    def sqTimeDensity(tdiff):
        weight = Distributions.PDFs.egamma(tdiff,3.75,1.5)
        return weight*weight

    @staticmethod
    def timeDensityNull(timediff):
        return Distributions.PDFs.expon(timediff,.015)

    @staticmethod
    def angDensityChase(angle1,angle2):
        angdif = math.fabs(angle1-angle2)
        bvar = angdif/math.pi/2.+.5
        """2pi*(angdif-.5) ~ Beta(2,2)"""
        weight = Distributions.PDFs.beta(bvar,2.,2.)
        return weight

    @staticmethod
    def angDensityNull(angle1,angle2):
        """null angle distribution is uniform on [0,2*pi]"""
        return -np.log(np.pi*2)

    @staticmethod
    def speedDensityChase(speed):
        """avg running speed ~14 mph when chased"""
        weight = Distributions.PDFs.gamma(speed,14.,1./.45)
        return weight

    @staticmethod
    def sqSpeedDensity(speed):
        weight = Distributions.PDFs.egamma(speed,14.,1./.45)
        return weight*weight

    @staticmethod
    def sqAngleDensity(angle):
        weight = Distributions.PDFs.ebeta(angle,2.,2.)
        return weight*weight

    @staticmethod
    def speedDensityNull(speed):
        """avg running speed ~9.4 mph"""
        weight = Distributions.PDFs.gamma(speed,6.,1/.7)
        return weight
